from enum import Enum
from typing import Optional, List, Dict, Any, Union, Tuple
from pydantic import BaseModel
from app.schemas.metrics_schema import Chart, BarChart, ScatterPlot, LineChart
from app.schemas.assets_schema import Text, Ontology, KnowledgeGraph, GraphPattern, GraphRule, MediaFile, TextRule


# 任务状态
class TaskStatus(str, Enum):
    PENDING = "PENDING"
    RUNNING = "RUNNING"
    COMPLETED = "COMPLETED"
    FAILED = "FAILED"


# MPLP算法的输入参数
class MPLPInputParams(BaseModel):
    nodes: List[Dict[str, Any]]  # 格式为 [{"id": id, "type": type, ...}, ...]
    edges: List[Dict[str, Any]]  # 格式为 [{"source": source_id, "target": target_id, ...}, ...]
    meta_paths: List[List[str]]  # 元路径定义，如 [["User", "Item", "User"], ["User", "Tag", "Item", "User"]]
    source_type: str  # 源节点类型
    target_type: str  # 目标节点类型
    top_k: int = 10  # 返回的预测数量


# MPLP算法的输出参数
class MPLPOutputParams(BaseModel):
    predictions: List[Dict[str, Any]]  # 预测结果：[{"source": source_id, "target": target_id, "score": score}, ...]
    algorithm: str = "MetaPathLinkPrediction"  # 算法名称
    parameters: Dict[str, Any]  # 算法参数
    path_counts: Dict[str, float]  # 元路径计数统计


# 输入参数
class InputParams(BaseModel):
    mplp_params: Optional[MPLPInputParams] = None


# 输出参数
class OutputParams(BaseModel):
    mplp_results: Optional[MPLPOutputParams] = None


# 算法请求
class AlgorithmRequest(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    input_params: InputParams


# 算法中间响应
class AlgorithmMiddleResponse(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    task_status: TaskStatus
    task_progress: int = 0
    task_logs: Optional[str] = None
    
    input_params: InputParams
    error_message: Optional[str] = None
    
    metrics: List[Chart] = []


# 算法最终响应
class AlgorithmResponse(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    task_status: TaskStatus
    task_progress: int = 100
    task_logs: Optional[str] = None
    
    error_message: Optional[str] = None
    
    output_params: OutputParams
    
    metrics: List[Chart] = [] 